package cn.bincker.mybatis.encrypt.data.migration.impl;

import cn.bincker.mybatis.encrypt.core.EncryptExecutor;
import cn.bincker.mybatis.encrypt.data.migration.MigrateDatabase;
import cn.bincker.mybatis.encrypt.data.migration.entity.ColumnInfo;
import cn.bincker.mybatis.encrypt.data.migration.entity.ForeignKeyInfo;
import cn.bincker.mybatis.encrypt.data.migration.entity.MigrateTable;
import cn.bincker.mybatis.encrypt.data.migration.exception.MigrateException;
import cn.bincker.mybatis.encrypt.entity.EncryptProperty;
import cn.bincker.mybatis.encrypt.entity.SecurityProperty;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.annotation.TableName;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import static cn.bincker.mybatis.encrypt.data.migration.utils.DateUtils.toHumanTime;
import static cn.bincker.mybatis.encrypt.data.migration.utils.NameUtils.*;

@Slf4j
@SuppressWarnings("SqlSourceToSinkFlow")
public abstract class AbstractMigrateDatabase implements MigrateDatabase {
    @Getter
    @Setter
    protected DataSource srcDataSource;
    @Getter
    @Setter
    protected DataSource distDataSource;
    @Getter
    @Setter
    private ProgressListener progressListener = (a, b, c, d)->{};
    @Getter
    @Setter
    private ExecutorService executorService;
    @Setter
    private int selectThreads;
    @Setter
    private int encryptThreads;
    @Setter
    private int insertThreads;
    private final Map<String, MigrateTable> migrateTableMap = new HashMap<>();
    @Getter
    @Setter
    private EncryptExecutor encryptExecutor;
    @Setter
    private Consumer<EncryptTask> onEncryptThreadInitListener;
    @Setter
    private Consumer<EncryptTask> onEncryptThreadDestroyListener;
    /**
     * 是否迁移所有表，这将会将没有注册的表也进行迁移，因为没有实体类，所以都是以原数据的形式迁移
     */
    private final boolean migrateAllTable;
    /**
     * 批量插入条数
     */
    private final int batchInsertRows;
    private final Map<MigrateTable, List<ColumnInfo>> columnsCache = new HashMap<>();
    private long prevSelectedRows = 0L;
    private long prevEncryptRows = 0L;
    private long prevMigratedRows = 0;
    private final Queue<InstantMigratedRows> perSecondMigratedRows = new LinkedList<>();
    private long prevTimestamp = 0L;
    private final LinkedBlockingQueue<Long> pageQueue = new LinkedBlockingQueue<>();
    private final LinkedBlockingQueue<Object[]> encryptQueue = new LinkedBlockingQueue<>();
    @Setter
    private int encryptQueueLimit = 100000;
    private final LinkedBlockingQueue<Object[]> insertQueue = new LinkedBlockingQueue<>();
    @Setter
    private int insertQueueLimit = 10000;
    private final List<SelectTask> selectTasks = new ArrayList<>();
    private final List<EncryptTask> encryptTasks = new ArrayList<>();
    private final List<InsertTask> insertTasks = new ArrayList<>();
    private final boolean failContinue;

    public AbstractMigrateDatabase(
            DataSource srcDataSource,
            DataSource distDataSource,
            EncryptExecutor encryptExecutor,
            boolean migrateAllTable,
            int batchInsertRows,
            boolean failContinue) {
        this.srcDataSource = srcDataSource;
        this.distDataSource = distDataSource;
        this.encryptExecutor = encryptExecutor;
        this.migrateAllTable = migrateAllTable;
        this.batchInsertRows = batchInsertRows;
        this.failContinue = failContinue;
        selectThreads = 1;
        encryptThreads = Runtime.getRuntime().availableProcessors() - 2;
        if (encryptThreads < 1) throw new MigrateException("too few processors are available. need 4 but got " + Runtime.getRuntime().availableProcessors());
        insertThreads = 1;
    }

    @Override
    public void registerTable(Class<?> entityClass) {
        registerTable(getMigrateTableByEntityClass(entityClass));
    }

    /**
     * 从Mybatis实体类获取表明，如果@TableName存在从TableName获取，不存在则将类名转换为表名
     */
    private static MigrateTable getMigrateTableByEntityClass(Class<?> entityClass) {
        var annotation = entityClass.getAnnotation(TableName.class);
        var result = new MigrateTable();
        result.setEntityClass(entityClass);
        if (annotation != null){
            result.setSrcName(annotation.value());
            result.setDistName(annotation.value());
            result.setSrcSchema(annotation.schema());
            result.setDistSchema(annotation.schema());
        }else{
            var tableName = entityClassToTableName(entityClass);
            result.setSrcName(tableName);
            result.setDistName(tableName);
        }
        return result;
    }

    @Override
    public void registerTable(MigrateTable table) {
        migrateTableMap.put(table.getSrcName(), table);
    }

    @Override
    public Collection<MigrateTable> getMigrateTables() {
        return migrateTableMap.values();
    }

    private static List<MigrateTable> getTables(Connection connection){
        var result = new ArrayList<MigrateTable>();
        DatabaseMetaData metaData;
        try {
            metaData = connection.getMetaData();
            var tables = metaData.getTables(null, null, null, new String[]{"TABLE"});
            while (tables.next()){
                MigrateTable table = new MigrateTable();
                var schema = tables.getString("TABLE_SCHEM");
                var tableName = tables.getString("TABLE_NAME");
                table.setSrcSchema(schema);
                table.setDistSchema(schema);
                table.setSrcName(tableName);
                table.setDistName(tableName);
                result.add(table);
            }
            return result;
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<ColumnInfo> getSrcColumns(Connection conn, MigrateTable table){
        return columnsCache.computeIfAbsent(
                table,
                t->getColumns(
                        conn,
                        table.getSrcSchema(),
                        table.getSrcName(),
                        table.getEntityClass() == null ? Collections.emptyMap() : encryptExecutor.getEncryptFields(table.getEntityClass())
                )
        );
    }

    @Override
    public List<ColumnInfo> getDistColumns(Connection conn, MigrateTable table) {
        return getColumns(
                conn,
                table.getDistSchema(),
                table.getDistName(),
                table.getEntityClass() == null ? Collections.emptyMap() : encryptExecutor.getEncryptFields(table.getEntityClass())
        );
    }

    private static List<ColumnInfo> getColumns(Connection connection, String schema, String table, Map<String, SecurityProperty> encryptPropertyMap){
        try {
            var metaData = connection.getMetaData();
            var primaryKeys = metaData.getPrimaryKeys(null, schema, table);
            var primaryKeysSet = new HashSet<String>();
            while (primaryKeys.next()) primaryKeysSet.add(primaryKeys.getString("COLUMN_NAME"));
            var foreignKeyResultSet = metaData.getImportedKeys(null, schema, table);
            var foreignKeys = new HashMap<String, ForeignKeyInfo>();
            while (foreignKeyResultSet.next()) {
                var fk = new ForeignKeyInfo(
                        foreignKeyResultSet.getString("FK_NAME"),
                        foreignKeyResultSet.getString("FKCOLUMN_NAME"),
                        foreignKeyResultSet.getString("PKTABLE_NAME"),
                        foreignKeyResultSet.getString("PKCOLUMN_NAME")
                );
                foreignKeys.put(fk.getColumn(), fk);
            }
            var columns = metaData.getColumns(null, schema, table, null);
            var result = new HashMap<String, ColumnInfo>();
            var index = 0;
            while (columns.next()){
                String columnName = columns.getString("COLUMN_NAME");
                var columnInfo = new ColumnInfo(
                        columnName,
                        columns.getString("TYPE_NAME"),
                        columns.getInt("COLUMN_SIZE"),
                        primaryKeysSet.contains(columnName),
                        "YES".equalsIgnoreCase(columns.getString("IS_NULLABLE")),
                        index++,
                        null,
                        foreignKeys.get(columnName)
                );
                result.put(columnName, columnInfo);
            }
            for (SecurityProperty property : encryptPropertyMap.values()) {
                if (!(property instanceof EncryptProperty)) continue;
                var encryptProperty = (EncryptProperty) property;
                var tableField = encryptProperty.getField() != null ? encryptProperty.getField().getAnnotation(TableField.class) : null;
                String fieldName;
                if (tableField != null){
                    fieldName = tableField.value();
                }else{
                    fieldName = camelToSnakeCase(encryptProperty.getName());
                }
                var column = result.get(fieldName);
                if (column == null) column = result.get(fieldName.toUpperCase());
                if (column == null) throw new MigrateException("加密字段[" + encryptProperty.getGetter() + "]不存在于表[" + schema + "." + table + "]中");
                column.setEncryptProperty(encryptProperty);
            }
            return result.values().stream().sorted(Comparator.comparingInt(ColumnInfo::getIndex)).collect(Collectors.toList());
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    private void verifyTable(MigrateTable table) {
        try (var srcConn = srcDataSource.getConnection(); var distConn = distDataSource.getConnection()){
            var srcColumns = getSrcColumns(srcConn, table);
            var distColumns = getDistColumns(distConn, table)
                    .stream()
                    .collect(Collectors.toMap(ColumnInfo::getName, c->c));
            var entityClass = table.getEntityClass();
            if (entityClass != null){
                for (ColumnInfo column : srcColumns) {
                    var dColumn = distColumns.get(column.getName());
                    if (dColumn == null){
                        try(var statement = distConn.createStatement()) {
                            String sql = buildAddEncryptColumnSql(table, column);
                            statement.execute(sql);
                        }
                    }
                    if(encryptExecutor.isEncryptField(entityClass, snakeCaseToCamel(column.getName()))){
                        checkEncryptColumn(distConn, table, dColumn);
                    }else{
                        if (!column.equals(dColumn)){
                            throw new RuntimeException("target table structure is incorrect, please drop the table: table=" + table.getDistTableName() + "\tcolumn=" + column.getName());
                        }
                    }
                }
            }else{
                for (ColumnInfo column : srcColumns) {
                    var dColumn = distColumns.get(column.getName());
                    if (dColumn != null){
                        if (!column.equals(dColumn)){
                            throw new RuntimeException("target table structure is incorrect, please drop the table: table=" + table.getDistTableName() + "\tcolumn=" + dColumn.getName());
                        }
                    }else{
                        throw new RuntimeException("target table structure is incorrect, please drop the table: table=" + table.getDistTableName() + "\tcolumn=" + column.getName());
                    }
                }
            }
        } catch (SQLException e) {
            throw new MigrateException(e);
        }
    }

    @Override
    public void migrate() {
        var migratedSuccessRows = new AtomicLong();
        var migratedFailedRows = new AtomicLong();
        var encryptedRows = new AtomicLong();
        var selectedRows = new AtomicLong();
        executorService = Executors.newFixedThreadPool(selectThreads + encryptThreads + insertThreads);
        for (int i = 0; i < selectThreads; i++) {
            SelectTask task = new SelectTask(
                    srcDataSource,
                    this::buildPageQuerySql,
                    this::buildExistsQuery,
                    pageQueue,
                    encryptQueue,
                    failContinue,
                    getPageSize(),
                    selectedRows,
                    migratedSuccessRows,
                    encryptQueueLimit
            );
            executorService.submit(task);
            selectTasks.add(task);
        }
        for (int j = 0; j < encryptThreads; j++) {
            EncryptTask task = new EncryptTask(
                    encryptExecutor,
                    encryptQueue,
                    insertQueue,
                    encryptedRows,
                    migratedFailedRows,
                    onEncryptThreadInitListener,
                    onEncryptThreadDestroyListener,
                    failContinue,
                    insertQueueLimit
            );
            executorService.submit(task);
            encryptTasks.add(task);
        }
        for (int i = 0; i < insertThreads; i++) {
            InsertTask task = new InsertTask(
                    distDataSource,
                    this::buildInsertSql,
                    this::beforeInsert,
                    insertQueue,
                    batchInsertRows,
                    migratedSuccessRows,
                    migratedFailedRows,
                    failContinue
            );
            executorService.submit(task);
            insertTasks.add(task);
        }
        try (var srcConn = srcDataSource.getConnection(); var distConn = distDataSource.getConnection()){
            if (migrateAllTable){
                for (var table : getTables(srcConn)) {
                    migrateTableMap.putIfAbsent(table.getSrcName(), table);
                }
            }
            var distTables = getTables(distConn);
            var tables = new ArrayList<>(migrateTableMap.values());
            try {
                for (int i = 0; i < tables.size(); i++) {
                    var startTime = System.currentTimeMillis();
                    var table = tables.get(i);
                    table.setColumns(getSrcColumns(srcConn, table));
                    if (distTables.stream().noneMatch(t->t.getSrcTableName().equals(table.getDistTableName()))) {
                        progressListener.onProgress(i + 1, tables.size(), table, "create table...");
                        try (var statement = distConn.createStatement()){
                            String ddl = createDistTableDDL(table);
                            log.info("create table: \n" + ddl);
                            if(!statement.execute(ddl)){
                                throw new MigrateException("create dist table fail: table=" + table.getDistTableName());
                            }
                        }
                    } else {
                        progressListener.onProgress(i + 1, tables.size(), table, "varify table...");
                        verifyTable(table);
                    }
                    encryptTasks.forEach(t->t.setMigrateTable(table));
                    insertTasks.forEach(t->t.setMigrateTable(table));
                    selectedRows.set(table.getMigratedRows().get());
                    prevSelectedRows = 0L;
                    prevTimestamp = System.currentTimeMillis();
                    migratedSuccessRows.set(table.getMigratedRows().get());
                    migratedFailedRows.set(0L);
                    encryptedRows.set(table.getMigratedRows().get());
                    prevEncryptRows = 0L;
                    boolean checkExist;
                    try(var statement = srcConn.createStatement(); var distStatement = distConn.createStatement()){
                        progressListener.onProgress(i + 1, tables.size(), table, "count src table...");
                        String sql = buildCountSql(table.getSrcTableSqlName());
                        log.info("count src rows: \n" + sql);
                        var result = statement.executeQuery(sql);
                        if(!result.next()) throw new MigrateException("select count fail: table=" + table.getSrcTableName());
                        table.setTotalRows(result.getLong(1));
                        progressListener.onProgress(i + 1, tables.size(), table, "count dist table...");
                        sql = buildCountSql(table.getDistTableSqlName());
                        log.info("count dist rows: \n" + sql);
                        result = distStatement.executeQuery(sql);
                        if(!result.next()) throw new MigrateException("select dist table count fail: table=" + table.getSrcTableName());
                        var distRows = result.getLong(1);
                        checkExist = distRows > 0;//是否进行判断数据是否存在
                        log.info("start encrypt: table={}\trows={}/{}\t{}", table.getSrcTableName(), distRows, table.getTotalRows(), Math.round(distRows * 1.0 / table.getTotalRows() * 10000) / 100 + "%");
                    }
                    selectTasks.forEach(t->t.setMigrateTable(table, checkExist));
                    for (
                            long page = table.getMigratedRows().get() / getPageSize(),
                            totalPage = table.getTotalRows() / getPageSize() + (table.getTotalRows() % getPageSize() == 0 ? 0 : 1);
                            page < totalPage;
                            page++
                    ) {
                        pageQueue.add(page);
                    }
                    while ((!pageQueue.isEmpty() || !encryptQueue.isEmpty() || !insertQueue.isEmpty()) && running() && table.getMigratedRows().get() < table.getTotalRows()){
                        try{
                            synchronized (table){
                                table.wait(1000);
                            }
                        }catch (InterruptedException e){
                            throw new MigrateException(e);
                        }
                        updateProgress(table.getMigratedRows(), selectedRows, encryptedRows, migratedFailedRows, i, tables, table);
                    }
                    if (!running()) break;
                    //继续等待让数据处理完
                    waitBusy();
                    progressListener.onProgress(
                            i + 1,
                            tables.size(),
                            table,
                            "finished, failed: " + migratedFailedRows.get() +
                                    "\ttook " + toHumanTime((int)((System.currentTimeMillis() - startTime)/1000))
                    );
                }
            }catch (Throwable e){
                throw new MigrateException(e);
            }
        } catch (SQLException e) {
            throw new MigrateException(e);
        } finally {
            selectTasks.forEach(SelectTask::close);
            encryptTasks.forEach(EncryptTask::close);
            insertTasks.forEach(InsertTask::close);
        }
        selectTasks.forEach(SelectTask::waitDestroyed);
        encryptTasks.forEach(EncryptTask::waitDestroyed);
        insertTasks.forEach(InsertTask::waitDestroyed);
        selectTasks.clear();
        encryptTasks.clear();
        insertTasks.clear();
        executorService.shutdown();
        try {
            if(!executorService.awaitTermination(30, TimeUnit.SECONDS)){
                log.warn("shutdown execute service fail.");
            }
        } catch (InterruptedException e) {
            throw new MigrateException(e);
        }
    }

    private void waitBusy() {
        selectTasks.forEach(SelectTask::waitBusy);
        encryptTasks.forEach(EncryptTask::waitBusy);
        insertTasks.forEach(InsertTask::waitBusy);
    }

    private boolean running() {
        return selectTasks.stream().allMatch(SelectTask::isRunning) && encryptTasks.stream().allMatch(EncryptTask::isRunning) && insertTasks.stream().allMatch(InsertTask::isRunning);
    }

    private void updateProgress(AtomicLong migratedRows, AtomicLong selectedRows, AtomicLong encryptedRows, AtomicLong failedRows, int i, ArrayList<MigrateTable> tables, MigrateTable table) {
        var totalRows = table.getTotalRows();
        var selectRows = selectedRows.get();
        var migrateRows = migratedRows.get();
        var encryptRows = encryptedRows.get();
        var timestamp = System.currentTimeMillis();
        var selectSpeed = (selectRows - prevSelectedRows) * 1.0 / ((timestamp - prevTimestamp) / 1000.0);
        var encryptSpeed = (encryptRows - prevEncryptRows) * 1.0 / ((timestamp - prevTimestamp) / 1000.0);
        var combinedSpeed = perSecondMigratedRows.isEmpty() ? 0 :
                perSecondMigratedRows.stream().mapToLong(InstantMigratedRows::getRows).sum() * 1.0 / ((timestamp - perSecondMigratedRows.peek().time) / 1000.0);
        var selectEta = selectSpeed <= 0 ? 0 : (totalRows - selectRows) / selectSpeed;
        var encryptEta = encryptSpeed <= 0 ? 0 : (totalRows - encryptRows) / encryptSpeed;
        var combinedEta = combinedSpeed <= 0 ? 0 : (totalRows - migrateRows) / combinedSpeed;
        progressListener.onProgress(
                i + 1,
                tables.size(),
                table,
                "select speed: " + ((int) selectSpeed) + "/s eta: " + toHumanTime((int) selectEta) +
                        "\tencrypt speed: " + ((int) encryptSpeed) + "/s eta: " + toHumanTime((int) encryptEta) +
                        "\tcombined speed: " + ((int) combinedSpeed) + "/s eta: " + toHumanTime((int) combinedEta) +
                        "\tfailed rows: " + failedRows.get()
        );
        perSecondMigratedRows.add(new InstantMigratedRows(migrateRows - prevMigratedRows, timestamp));
        if (perSecondMigratedRows.size() > 10) perSecondMigratedRows.poll();
        prevSelectedRows = selectRows;
        prevEncryptRows = encryptRows;
        prevMigratedRows = migrateRows;
        prevTimestamp = timestamp;
    }

    protected abstract void beforeInsert(Connection conn, MigrateTable table);

    /**
     * 构建表插入语句
     * @param table 表名
     */
    protected abstract String buildInsertSql(MigrateTable table);

    /**
     * 构建数据是否存在(已同步)查询
     * @param connection 同步目标数据库连接
     * @param table 表信息
     */
    protected abstract PreparedStatement buildExistsQuery(Connection connection, MigrateTable table, List<Object[]> rows) throws SQLException;

    protected abstract String buildCountSql(String tableName);

    protected abstract String buildPageQuerySql(MigrateTable table, long page);

    protected abstract String createDistTableDDL(MigrateTable table);

    protected abstract int getPageSize();

    protected abstract void checkEncryptColumn(Connection connection, MigrateTable table, ColumnInfo column);

    protected abstract String buildAddEncryptColumnSql(MigrateTable table, ColumnInfo column);

    @Data
    @AllArgsConstructor
    static class InstantMigratedRows {
        private final long rows;
        private final long time;
    }
}
